#ifndef _MEM_HOOK_H_
#define _MEM_HOOK_H_

#include <stdint.h>
#include <stddef.h>

#include <climits>
#include <map>
#include <mutex>

typedef void* (*MallocFn)(size_t);
typedef void* (*FreeFn)(void*);

extern MallocFn real_malloc;
extern FreeFn real_free;

template <class T>
inline T* _allocate(ptrdiff_t size, T*)
{
	/*std::cout << "_allocate called" << std::endl;
	set_new_handler(0);
	T* tmp = (T*)(::operator new((size_t)(size * sizeof(T))));
	if (NULL == tmp) {
	    std::cerr << "out of memory" << std::endl;
	    exit(0);
	}
	*/

	void* p = real_malloc((size_t)(size * sizeof(T)));
	if (!p) {
		throw std::bad_alloc();
	}
	return static_cast<T*>(p);
}

template <class T>
inline void _deallocate(T* p)
{
	//::operator delete(p);
	real_free(p);
}

template <class T1, class T2>
inline void _construct(T1* p, const T2& value)
{
	::new (p) T1(value);
}

template <class T>
inline void _destroy(T* p)
{
	p->~T();
}

template <typename T>
struct CustomAllocator {
	using value_type = T;
	using pointer = T *;
	using const_pointer = const T *;
	using reference = T &;
	using const_reference = const T &;
	using size_type = size_t;
	using difference_type = ptrdiff_t;

	template <typename U>
	struct rebind {
		using other = CustomAllocator<U>;
	};

	pointer allocate(size_type n, const void* hint = 0)
	{
		return _allocate(static_cast<difference_type>(n), static_cast<pointer>(nullptr));
	}

	void deallocate(pointer p, size_type n)
	{
		_deallocate(p);
	}

	void construct(pointer p, const T& value)
	{
		_constructor(p, value);
	}

	void destroy(pointer p)
	{
		_destroy(p);
	}

	pointer address(reference x)
	{
		return static_cast<pointer>(&x);
	}

	const_pointer const_address(const_reference x)
	{
		return static_cast<const_pointer>(&x);
	}

	size_type max_size() const
	{
		return static_cast<size_type>(UINT_MAX / sizeof(T));
	}
};

class MemHook {
public:
	MemHook();
	~MemHook();

	void* Malloc(size_t n);

	void Free(void* p);

	int Init();

private:
	typedef std::map<void*, size_t, std::less<void*>,
	        CustomAllocator<std::pair<void* const, size_t>>>  AllocatorMap;

	AllocatorMap* allocate_ptrs;
	std::mutex*   mutex;
	bool         initialized;
};

#endif
